import sensor, image, lcd, time
import KPU as kpu
import gc, sys
from machine import UART
from fpioa_manager import fm
from machine import WDT, reset
import machine


input_size = (224, 224)
labels = ['watermelon', 'liulian', 'banana', 'apple']

def lcd_show_except(e):
    import uio
    err_str = uio.StringIO()
    sys.print_exception(e, err_str)
    err_str = err_str.getvalue()
    img = image.Image(size=input_size)
    img.draw_string(0, 10, err_str, scale=1, color=(0xff,0x00,0x00))
    lcd.display(img)

class Comm:
    def __init__(self, uart):
        self.uart = uart
        self.receive_buffer = bytearray()      #接收缓冲区

    def send_ready(self):
        self.uart.write("[SYSTEM] READY\n".encode())

    def check_command(self):
        max_read_size = 128
        while self.uart.any():
            data = self.uart.read(max_read_size)
            if data:
                self.receive_buffer.extend(data)
            else:
                break

        # 手动查找 b'1' 的位置
        target = ord('1')  # 获取字符 '1' 的 ASCII 值（49）
        found_index = -1
        for i in range(len(self.receive_buffer)):
            if self.receive_buffer[i] == target:
                found_index = i
                break

        if found_index != -1:
            # 保留未处理的数据（从 found_index+1 开始）
            self.receive_buffer = self.receive_buffer[found_index+1:]
            return True
        return False

    def send_classify_result(self, pmax, idx, label):
        msg = "{}:{:.2f}:{}\n".format(idx, pmax, label)
        self.uart.write(msg.encode())

def init_uart():
    fm.register(10, fm.fpioa.UART1_TX, force=True)
    fm.register(11, fm.fpioa.UART1_RX, force=True)

    uart = UART(UART.UART1, 115200, 8, 0, 1, timeout=1000, read_buf_len=256)
    return uart

def main(labels = None, model_addr="/sd/m.kmodel", sensor_window=input_size, lcd_rotation=0, sensor_hmirror=True, sensor_vflip=True):
    #初始化看门狗（主循环看门狗）
    main_wdt = WDT(timeout=10000)  #5秒超时

    #硬件初始化
    sensor.reset()
    sensor.set_pixformat(sensor.RGB565)
    sensor.set_framesize(sensor.QVGA)
    sensor.set_windowing(sensor_window)
    sensor.set_hmirror(sensor_hmirror)
    sensor.set_vflip(sensor_vflip)
    sensor.run(1)

    lcd.init(type=1)
    lcd.rotation(lcd_rotation)
    lcd.clear(lcd.WHITE)

    #加标签
    if not labels:
        try:
            with open('labels.txt','r') as f:
                exec(f.read())
        except:
            pass
    if not labels:
        print("no labels.txt")
        img = image.Image(size=(320, 240))
        img.draw_string(90, 110, "no labels.txt", color=(255, 0, 0), scale=2)
        lcd.display(img)
        time.sleep(2)
        return 1

    #显示启动画面
    try:
        img = image.Image("startup.jpg")
        lcd.display(img)
    except Exception:
        img = image.Image(size=(320, 240))
        img.draw_string(90, 110, "loading model...", color=(255, 255, 255), scale=2)
        lcd.display(img)

    #初始化 串口
    uart = init_uart()
    comm = Comm(uart)

    task = None

    try:
        #加载模型
        task = kpu.load(model_addr)
        print("Model loaded successfully")

        # 发送系统就绪信号
        comm.send_ready()
        print("[SYSTEM] Ready signal sent")

        while(True):        #添加内存回收，防止未及时回收内存，导致内存耗尽
            main_wdt.feed()  #喂主循环看门狗

            #图像处理
            img = sensor.snapshot()

            # 新增：检查串口指令--------------------------------
            if comm.check_command():
                print("[CMD] Received classify command")
                cmd_time = time.ticks_ms()
                fmap = kpu.forward(task, img)
                cmd_time = time.ticks_ms() - cmd_time

                plist = fmap[:]
                pmax = max(plist)
                max_index = plist.index(pmax)

                # 立即发送结果不等待显示刷新
                comm.send_classify_result(pmax, max_index, labels[max_index].strip())
                print("[CMD] Force classification sent")
                continue  # 跳过本次循环的显示步骤
            # ------------------------------------------------

            #模型推理
            start_time = time.ticks_ms()
            fmap = kpu.forward(task, img)     #fmap是模型的输出
            inference_time = time.ticks_ms() - start_time

            #解析推理
            plist=fmap[:]
            pmax=max(plist)
            max_index=plist.index(pmax)

            #显示结果
            img.draw_string(0,0, "%.2f : %s" %(pmax, labels[max_index].strip()), scale=2, color=(255, 0, 0))
            img.draw_string(0, 200, "t:%dms" %(inference_time), scale=2, color=(255, 0, 0))
            #如果模型返回的结果格式不符合预期，调用comm.send_classify_result()

            #串口发送
            #comm.send_classify_result(pmax, max_index, labels[max_index].strip())

            #刷新显示
            lcd.display(img)
            gc.collect()   #每次循环结束后强制回收内存
    #新增位置3：显示捕获IDE中断信号
    except KeyboardInterrupt as e:
        print("\n用户手动停止脚本")
        main_wdt.stop()
        return 0   #返回特定状态码
    except Exception as e:
        print("model loaded failed:",e)
        main_wdt.stop()
        raise
    finally:
        #新增位置四，严格按顺序释放所有硬件资源
        print("进入finally块")

        #资源释放顺序优化
        if task is not None:
            kpu.deinit(task)
            print("模型资源已释放")
        sensor.shutdown(True)     #关闭摄像头硬件
        print("摄像头已关闭")
        lcd.clear()      #清空屏幕残留图像
        lcd.deinit()     #关闭LCD控制器
        print("LCD已关闭")

        #关闭串口
        uart.deinit()
        print("硬件外设已全部关闭")
        main_wdt.stop()   #停止看门狗
        print("WDT stopped")
        time.sleep(1)    #等待硬件状态稳定
        #reset()   #使用硬件复位确保状态干净
        import utime
        utime.sleep(1)


if __name__ == "__main__":
    global_wdt = WDT(timeout=30000)
    try:
        ret = main(labels=labels, model_addr="/sd/model-98702.kmodel")
        if ret == 0:
            print("程序正常退出")
    except KeyboardInterrupt:
        print("\n[DEBUG] 用户强制终止，停止外层看门狗")
        global_wdt.stop()
    except Exception as e:
        sys.print_exception(e)
        lcd_show_except(e)
    finally:
        global_wdt.stop()  # 确保外层看门狗停止
        print("[DEBUG] 外层看门狗已强制停止")
